#ifndef _VULKAN_HELPER_H_
#define _VULKAN_HELPER_H_

#include "std.hpp"

namespace Vulkan {
	uint32_t getMemoryType(uint32_t typeBits, VkMemoryPropertyFlags properties)
	{
		for (uint32_t i = 0; i < context.memoryPropertes.memoryTypeCount; i++)
		{
			if ((typeBits & 1) == 1)
			{
				if ((context.memoryPropertes.memoryTypes[i].propertyFlags & properties) == properties)
				{
					return i;
				}
			}
			typeBits >>= 1;
		}
		throw std::runtime_error("Could not find a matching memory type");
	}

	VkDeviceMemory allocateMemory(const VkDeviceSize size, uint32_t memoryTypeIndex) {
		VkMemoryAllocateInfo mem_alloc = {};
		mem_alloc.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
		mem_alloc.allocationSize = size;
		mem_alloc.memoryTypeIndex = memoryTypeIndex;

		VkDeviceMemory memory = VK_NULL_HANDLE;
		vkAllocateMemory(context.device, &mem_alloc, nullptr, &memory);
		return memory;
	}

	VkDeviceMemory createMemory(const VkImage& image, VkMemoryPropertyFlags properties) {
		VkMemoryRequirements memReqs;
		vkGetImageMemoryRequirements(context.device, image, &memReqs);
		auto memory = allocateMemory(memReqs.size,getMemoryType(memReqs.memoryTypeBits, properties));
		vkBindImageMemory(context.device, image, memory, 0);
		return memory;
	}

	VkDeviceMemory createMemory(const VkBuffer& buffer, VkMemoryPropertyFlags properties) {
		VkMemoryRequirements memReqs;
		vkGetBufferMemoryRequirements(context.device,buffer,&memReqs);
		auto memory = allocateMemory(memReqs.size,getMemoryType(memReqs.memoryTypeBits, properties));
		vkBindBufferMemory(context.device,buffer,memory,0);
		return memory;
	}

	VkBuffer createBuffer(VkDeviceSize size,VkBufferUsageFlags usage){
		VkBufferCreateInfo info = {};
		info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
		info.size = size;
		info.usage = usage;
		info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
		VkBuffer buffer = VK_NULL_HANDLE;
		vkCreateBuffer(context.device, &info, nullptr, &buffer);
		return buffer;
	}

	VkDescriptorPool createDescriptorPool(std::vector<VkDescriptorPoolSize>& poolSizes,uint32_t maxSets) {
		VkDescriptorPoolCreateInfo info = {};
		info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
		info.poolSizeCount = (uint32_t)poolSizes.size();
		info.pPoolSizes = poolSizes.data();
		info.maxSets = maxSets;
		VkDescriptorPool descriptorPool = VK_NULL_HANDLE;
		vkCreateDescriptorPool(context.device, &info, nullptr, &descriptorPool);
		return descriptorPool;
	}

	VkDescriptorSet createDescriptorSet(VkDescriptorPool descriptorPool,const VkDescriptorSetLayout* pSetLayouts){
		VkDescriptorSetAllocateInfo info = {};
		info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
		info.descriptorPool = descriptorPool;
		info.pSetLayouts = pSetLayouts;
		info.descriptorSetCount = 1;
		VkDescriptorSet set = VK_NULL_HANDLE;
		vkAllocateDescriptorSets(context.device, &info, &set);
		return set;
	}

	inline const std::string getResourcePath(const std::string& path) {
		return std::string(SOURCE_PATH) + "/" + path;
	}

	inline const std::vector<char> readShaderFile(const std::string& path) {
		std::ifstream file(path, std::ios::ate | std::ios::binary);

		if (!file.is_open()) {
			throw std::runtime_error("failed to open file!");
		}
		size_t fileSize = (size_t)file.tellg();
		std::vector<char> buffer(fileSize);
		file.seekg(0);
		file.read(buffer.data(), fileSize);
		file.close();
		return buffer;
	}

	inline VkShaderModule createShaderModule(const std::string& path) {
		auto code = readShaderFile(path);
		VkShaderModule module = VK_NULL_HANDLE;
		VkShaderModuleCreateInfo info = {};
		info.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
		info.codeSize = code.size();
		info.pCode = (uint32_t*)code.data();
		vkCreateShaderModule(context.device, &info, nullptr, &module);
		return module;
	}

	template <typename T>
	void mapData(VkDeviceMemory memory,VkDeviceSize size,T* data){
		T* pointer;
		vkMapMemory(context.device, memory, 0, size, 0, (void **)&pointer);
		memcpy(pointer, data, size);
		vkUnmapMemory(context.device, memory);
	}

	std::vector<VkCommandBuffer> createCommandBuffer(uint32_t count){
		VkCommandBufferAllocateInfo info = {};
		info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
		info.commandPool = context.commandPool;
		info.commandBufferCount = count;
		info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
		std::vector<VkCommandBuffer> buffers(count,VK_NULL_HANDLE);
		vkAllocateCommandBuffers(context.device,&info,buffers.data());
		return buffers;
	}

	VkDescriptorSetLayout createSetLayout(std::vector<VkDescriptorSetLayoutBinding>& bindings) {
		VkDescriptorSetLayoutCreateInfo info = {};
		info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
		info.pBindings = bindings.data();
		info.bindingCount = (uint32_t)bindings.size();
		VkDescriptorSetLayout setLayout = VK_NULL_HANDLE;
		vkCreateDescriptorSetLayout(context.device, &info, nullptr, &setLayout);
		return setLayout;
	}

	VkCommandBuffer beginSingleCommandBuffer(){
		VkCommandBufferBeginInfo info = {};
		info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
		info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
		auto cmd = createCommandBuffer(1)[0];
		vkBeginCommandBuffer(cmd,&info);
		return cmd;
	}

	void flushSingleCommandBuffer(VkCommandBuffer cmd){
		vkEndCommandBuffer(cmd);
		VkSubmitInfo submitInfo = {};
		submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
		submitInfo.commandBufferCount = 1;
		submitInfo.pCommandBuffers = &cmd;
		vkQueueSubmit(context.graphicsQueue, 1, &submitInfo, VK_NULL_HANDLE);
		vkQueueWaitIdle(context.graphicsQueue);
		vkFreeCommandBuffers(context.device,context.commandPool,1,&cmd);
		cmd = VK_NULL_HANDLE;
	}

	void createBufferMemoryWithStaging(uint32_t size, void* data, VkBuffer& buffer, VkDeviceMemory& memory, VkBufferUsageFlags usage) {
		VkBuffer stagingBuffer = VK_NULL_HANDLE;
		VkDeviceMemory stagingMemory = VK_NULL_HANDLE;
		stagingBuffer = createBuffer(size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT);
		stagingMemory = createMemory(stagingBuffer, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
		mapData(stagingMemory, size, data);
		buffer = createBuffer(size, VK_BUFFER_USAGE_TRANSFER_DST_BIT | usage);
		memory = createMemory(buffer, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

		VkBufferCopy copyRegion = { 0,0,size };
		auto cmd = beginSingleCommandBuffer();
		vkCmdCopyBuffer(cmd, stagingBuffer, buffer, 1, &copyRegion);
		flushSingleCommandBuffer(cmd);

		vkDestroyBuffer(context.device, stagingBuffer, nullptr);
		vkFreeMemory(context.device, stagingMemory, nullptr);
	}

	VkSemaphore createSemaphore(){
		VkSemaphoreCreateInfo info = {};
		info.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;
		VkSemaphore semaphore = VK_NULL_HANDLE;
		vkCreateSemaphore(context.device,&info, nullptr,&semaphore);
		return semaphore;
	}

	VkPipelineLayout createPipelineLayout(uint32_t setLayoutCount, VkDescriptorSetLayout* setLayout) {
		VkPipelineLayoutCreateInfo info = {};
		info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
		info.setLayoutCount = setLayoutCount;
		info.pSetLayouts = setLayout;
		VkPipelineLayout layout = VK_NULL_HANDLE;
		vkCreatePipelineLayout(context.device, &info, nullptr, &layout);
		return layout;
	}

	void setImageLayout(
			VkCommandBuffer cmdbuffer,
			VkImage image,
			VkImageAspectFlags aspectMask,
			VkImageLayout oldImageLayout,
			VkImageLayout newImageLayout,
			VkImageSubresourceRange subresourceRange)
	{
		// Create an image barrier object
		VkImageMemoryBarrier imageMemoryBarrier = {};
		imageMemoryBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
		imageMemoryBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
		imageMemoryBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
		imageMemoryBarrier.oldLayout = oldImageLayout;
		imageMemoryBarrier.newLayout = newImageLayout;
		imageMemoryBarrier.image = image;
		imageMemoryBarrier.subresourceRange = subresourceRange;
		switch (oldImageLayout)
		{
			case VK_IMAGE_LAYOUT_UNDEFINED:
				// Image layout is undefined (or does not matter)
				// Only valid as initial layout
				// No flags required, listed only for completeness
				imageMemoryBarrier.srcAccessMask = 0;
				break;

			case VK_IMAGE_LAYOUT_PREINITIALIZED:
				// Image is preinitialized
				// Only valid as initial layout for linear images, preserves memory contents
				// Make sure host writes have been finished
				imageMemoryBarrier.srcAccessMask = VK_ACCESS_HOST_WRITE_BIT;
				break;

			case VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL:
				// Image is a color attachment
				// Make sure any writes to the color buffer have been finished
				imageMemoryBarrier.srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
				break;

			case VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL:
				// Image is a depth/stencil attachment
				// Make sure any writes to the depth/stencil buffer have been finished
				imageMemoryBarrier.srcAccessMask = VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
				break;

			case VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL:
				// Image is a transfer source
				// Make sure any reads from the image have been finished
				imageMemoryBarrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
				break;

			case VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL:
				// Image is a transfer destination
				// Make sure any writes to the image have been finished
				imageMemoryBarrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
				break;

			case VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL:
				// Image is read by a shader
				// Make sure any shader reads from the image have been finished
				imageMemoryBarrier.srcAccessMask = VK_ACCESS_SHADER_READ_BIT;
				break;
		}

		// Target layouts (new)
		// Destination access mask controls the dependency for the new image layout
		switch (newImageLayout)
		{
			case VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL:
				// Image will be used as a transfer destination
				// Make sure any writes to the image have been finished
				imageMemoryBarrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
				break;

			case VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL:
				// Image will be used as a transfer source
				// Make sure any reads from and writes to the image have been finished
				imageMemoryBarrier.srcAccessMask = imageMemoryBarrier.srcAccessMask | VK_ACCESS_TRANSFER_READ_BIT;
				imageMemoryBarrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
				break;

			case VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL:
				// Image will be used as a color attachment
				// Make sure any writes to the color buffer have been finished
				imageMemoryBarrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
				imageMemoryBarrier.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
				break;

			case VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL:
				// Image layout will be used as a depth/stencil attachment
				// Make sure any writes to depth/stencil buffer have been finished
				imageMemoryBarrier.dstAccessMask = imageMemoryBarrier.dstAccessMask | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
				break;

			case VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL:
				// Image will be read in a shader (sampler, input attachment)
				// Make sure any writes to the image have been finished
				if (imageMemoryBarrier.srcAccessMask == 0)
				{
					imageMemoryBarrier.srcAccessMask = VK_ACCESS_HOST_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT;
				}
				imageMemoryBarrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
				break;
		}

		// Put barrier on top
		VkPipelineStageFlags srcStageFlags = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
		VkPipelineStageFlags destStageFlags = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;

		// Put barrier inside setup command buffer
		vkCmdPipelineBarrier(
				cmdbuffer,
				srcStageFlags,
				destStageFlags,
				0,
				0, nullptr,
				0, nullptr,
				1, &imageMemoryBarrier);
	}
}

#endif